import copy
import numpy as np


class TD:
    """"
    n步TD,只包含算法，无法单独运行
    其中n=1为TD(0)

    eval:策略评估
    train:策略改进
    """
    def __init__(self,env,pie,select_A,V,Q,n=1,gamma=1) -> None:
        ## 策略
        self.pie=pie

        ## n步时序差分
        self.n=n

        ## 
        self.gamma=gamma
     
        ## 环境(包含step,sample,reset等关键函数)
        self.env=env
       


        ## 价值函数
        self.V=V

        ## Q函数
        self.Q=Q

        ## 动作选择函数
        self.select_A=select_A

        self.loss=[]

    def transform_act_index(self,a,index=True):
        """
        如果index为False,则输出实际动作
        否则，输出动作对应序号
        """
        if index:
            return np.where(np.array(self.env.action_space)==a)[0][0]

        return  self.env.action_space[int(a)]


    def eval(self,alpha=0.01,epoch=100):
        """
        时序差分策略评估

        """
        for _ in range(epoch):
            T=1e+8
            t=0
            s=self.env.reset()
            s_list=[s]
            r_list=[]
            # a=self.sample()
            while True:
                if t<T:
                    ## 生成下一幕s和r以及是否是终止点
                    a=self.select_A(s,self.env.action_space,pie=self.pie)
                    
                    ## 生成res-->s(t+1),r(t+1),flag,infos
                    res=self.env.step(a)
                    s=res[0]
                    s_list.append(s)
                    r_list.append(res[1])

                    if res[2]:
                        T=t+1
                    

                ## 步数+1
                t+=1

                ## 计算是否进入更新状态
                idx=t-self.n
                G=0
                if idx>=0:
                    
                    ## 计算回报
                    for i in range(min(self.n,T-idx)):
                        print("R({})+".format(idx+i+1),end="")
                        G+=(self.gamma**i)*r_list[idx+i]

                    ## 计算截断回报
                    if idx+self.n<T:
                        print("V({})".format(idx+self.n))
                        G+=(self.gamma**self.n)*self.V[s_list[idx+self.n]]

                    else:
                        G+=0

                    print(" ")
                    value=copy.deepcopy(self.V[s_list[idx]])
                    ## 更新状态函数
                    self.V[s_list[idx]]=value+alpha*(G-value)

                ## 在终止时刻前一刻定制
                if idx>=T-1:
                    break
        

        
                
                        


    def train(self,epoch=100,alpha=0.1,update=None):
        """
        n步差分改进策略
        n步SARSA 
        """
        for _ in range(epoch):
            T=1e+8
            t=0
            s=self.env.reset()
            a=self.env.sample()
            s_list=[s]
            r_list=[]
            a_list=[a]
            while True:
                if t<T:
                    ## 生成下一幕s和r以及是否是终止点
                    res=self.env.step(a)
                    s=res[0]
                    s_list.append(s)
                    r_list.append(res[1])
                    

                    if res[2]:
                        T=t+1

                    else:
                        ## 在幕结束之前生成动作A
                        a=self.select_A(s,self.env.action_space,self.pie)
                        a_list.append(a)
                    

                    
                ## 幕数+1
                t+=1
                idx=t-self.n
                G=0
                if idx>=0:
                    
                    for i in range(min(self.n,T-idx)):
                        G+=(self.gamma**i)*r_list[idx+i]

                    if idx+self.n<T:
                        # print("动作{}".format(a_list[idx+self.n]))
                        act_idx=self.transform_act_index(a_list[idx+self.n])

                        # print("动作序号{}".format(act_idx))
                        G+=(self.gamma**self.n)*self.Q[s_list[idx+self.n]][act_idx]

                    else:
                        G+=0

                    act_idx=self.transform_act_index(a_list[idx])
                    # print("更新状态:{},动作:{},收益:{}".format(s_list[idx],a_list[idx],G))
                    value=copy.deepcopy(self.Q[s_list[idx]][act_idx])
                    self.Q[s_list[idx]][act_idx]=value+alpha*(G-value)
                    if update:
                        self.pie=update(pie=self.pie,Q=self.Q)

                if idx>=T-1:
                        break



        